import numpy as np

import pyqtgraph as pg
import pyqtgraph.opengl as gl
from pyqtgraph import functions as fn
from pyqtgraph.Qt import QtCore
''''''
from robotdriver import robot_driver, my_robot
theta_1 = 0.0
theta_2 = -np.pi/2
theta_3 = 0.0
theta_4 = 0.0
theta_5 = 0.0
theta_6 = 0.0
joints = [theta_1, theta_2, theta_3, theta_4, theta_5, theta_6]
robot_driver.move_to_home(joints)
if robot_driver.is_ready():
    joints = robot_driver.get_joints_status()
    # print("show: robot_joints=", joints)
    robot_tf_tree = my_robot.get_robot_tf_tree(joints)
    ### base_box
    tf_matrix_base_0 = my_robot.get_base_robot_tf()
    base_origin = (0, 0, 0)
    quat_base = (0, 0, 0, 1)
    ### link 1
    tf_matrix_base_1 = np.dot(robot_tf_tree[0], tf_matrix_base_0)
    link1_origin = (tf_matrix_base_1[0,3], tf_matrix_base_1[1,3], tf_matrix_base_1[2,3])
    ### link 2
    tf_matrix_base_2 = np.dot(tf_matrix_base_1, robot_tf_tree[1])
    link2_origin = (tf_matrix_base_2[0,3], tf_matrix_base_2[1,3], tf_matrix_base_2[2,3])
    ### link 3
    tf_matrix_base_3 = np.dot(tf_matrix_base_2, robot_tf_tree[2])
    link3_origin = (tf_matrix_base_3[0,3], tf_matrix_base_3[1,3], tf_matrix_base_3[2,3])
    ### link 4
    tf_matrix_base_4 = np.dot(tf_matrix_base_3, robot_tf_tree[3])
    link4_origin = (tf_matrix_base_4[0,3], tf_matrix_base_4[1,3], tf_matrix_base_4[2,3])
    ### link 5
    tf_matrix_base_5 = np.dot(tf_matrix_base_4, robot_tf_tree[4])
    link5_origin = (tf_matrix_base_5[0,3], tf_matrix_base_5[1,3], tf_matrix_base_5[2,3])
    ### link 6
    tf_matrix_base_6 = np.dot(tf_matrix_base_5, robot_tf_tree[5])
    link6_origin = (tf_matrix_base_6[0,3], tf_matrix_base_6[1,3], tf_matrix_base_6[2,3])
    ### TCP
    tf_matrix_base_tcp = np.dot(tf_matrix_base_6, my_robot.get_end_tcp_tf())
    tcp_origin = (tf_matrix_base_tcp[0,3], tf_matrix_base_tcp[1,3], tf_matrix_base_tcp[2,3])
    ##########
''''''
app = pg.mkQApp("MY 6-DOF RoboticArm")
w = gl.GLViewWidget()
w.show()
w.setWindowTitle('MY 6-DOF RoboticArm')
w.setCameraPosition(distance=1200, azimuth=-45)

gx = gl.GLGridItem()
gx.setSize(1000, 1000, 1)
gx.setSpacing(50, 50, 1)
w.addItem(gx)

axisitem = gl.GLAxisItem()
axisitem.setSize(x=300, y=300, z=300)
w.addItem(axisitem)

txt_x = gl.GLTextItem()
txt_x.setData(pos=(300, 0, 0), color=(127, 255, 127, 255), text='X')
w.addItem(txt_x)

txt_y = gl.GLTextItem()
txt_y.setData(pos=(0, 300, 0), color=(127, 255, 127, 255), text='Y')
w.addItem(txt_y)

txt_z = gl.GLTextItem()
txt_z.setData(pos=(0, 0, 300), color=(127, 255, 127, 255), text='Z')
w.addItem(txt_z)
########################################################################################
def get_link_points(tf_matrix):
    point_e1_ = np.array([0.0, 0.0, 20.0, 1.0])
    point_e2_ = np.array([0.0, 0.0, -20.0, 1.0])
    point_e1 = (np.dot(tf_matrix, point_e1_))[0:3]
    point_e2 = (np.dot(tf_matrix, point_e2_))[0:3]
    return np.array([point_e1, point_e2])
    
pos = np.array([base_origin, 
                link1_origin, 
                link2_origin, 
                link3_origin, 
                link4_origin, 
                link5_origin, 
                link6_origin, 
                tcp_origin])
size = np.array([10, 10, 10, 10, 10, 10, 10, 10])
color = np.array([(1.0, 1.0, 0.0, 0.5),
                    (1.0, 1.0, 0.0, 0.5),
                    (1.0, 1.0, 0.0, 0.5),
                    (1.0, 1.0, 0.0, 0.5),
                    (1.0, 1.0, 0.0, 0.5),
                    (1.0, 1.0, 0.0, 0.5),
                    (1.0, 1.0, 0.0, 0.5),
                    (1.0, 1.0, 0.0, 0.5)])
sp_robot_point = gl.GLScatterPlotItem(pos=pos, size=size, color=color, pxMode=False)
w.addItem(sp_robot_point)
color = np.array([(1.0, 1.0, 1.0, 1),
                    (1.0, 1.0, 1.0, 1),
                    (1.0, 1.0, 1.0, 1),
                    (1.0, 1.0, 1.0, 1),
                    (1.0, 1.0, 1.0, 1),
                    (1.0, 1.0, 1.0, 1),
                    (1.0, 1.0, 1.0, 1),
                    (1.0, 1.0, 1.0, 1)])
lp_line = gl.GLLinePlotItem(pos=pos, color=color, width=5, antialias=True)
w.addItem(lp_line)
##
color = np.array([(0, 191/255.0, 1, 1), (0, 191/255.0, 1, 1)])
pos_link1 = get_link_points(tf_matrix_base_1)
lp_link1 = gl.GLLinePlotItem(pos=pos_link1, color=color, width=80, antialias=False)
pos_link2 = get_link_points(tf_matrix_base_2)
lp_link2 = gl.GLLinePlotItem(pos=pos_link2, color=color, width=20, antialias=True)
pos_link3 = get_link_points(tf_matrix_base_3)
lp_link3 = gl.GLLinePlotItem(pos=pos_link3, color=color, width=20, antialias=True)
pos_link4 = get_link_points(tf_matrix_base_4)
lp_link4 = gl.GLLinePlotItem(pos=pos_link4, color=color, width=20, antialias=True)
pos_link5 = get_link_points(tf_matrix_base_5)
lp_link5 = gl.GLLinePlotItem(pos=pos_link5, color=color, width=20, antialias=True)
pos_link6 = get_link_points(tf_matrix_base_6)
lp_link6 = gl.GLLinePlotItem(pos=pos_link6, color=color, width=20, antialias=True)
w.addItem(lp_link1)
w.addItem(lp_link2)
w.addItem(lp_link3)
w.addItem(lp_link4)
w.addItem(lp_link5)
w.addItem(lp_link6)

if __name__ == '__main__':
    pg.exec()
